{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### offline_test 脚本\n",
    "这个book脚本，设计目的是输出sam的分割结果。注意，分割的数据是无标注的数据集。需要人工验证分割效果。\n",
    "这里我们认为，sam的分割效果仍然强于base模型。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mmseg.apis import init_segmentor, inference_segmentor, show_result_pyplot\n",
    "from mmseg.core.evaluation import get_palette\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import cv2\n",
    "from sklearn.cluster import DBSCAN\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import os.path as osp\n",
    "import glob\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 初始模型加载\n",
    "# 不加载弱分割器\n",
    "# 强base分割器加载\n",
    "config_file = '/home/yangshuo/past_comp/DPLBV3P/code/1_SAM/weight/aug_all.py'\n",
    "checkpoint_file = '/home/yangshuo/past_comp/DPLBV3P/code/1_SAM/weight/iter_40000.pth'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "load checkpoint from local path: /home/yangshuo/past_comp/DPLBV3P/code/1_SAM/weight/iter_40000.pth\n"
     ]
    }
   ],
   "source": [
    "# 初始模型初始化\n",
    "model = init_segmentor(config_file, checkpoint_file, device='cuda:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sam的加载和初始化\n",
    "sam_checkpoint = \"/home/yangshuo/past_comp/segment-anything/pretrain/sam_vit_h_4b8939.pth\"\n",
    "device = \"cuda:1\"\n",
    "model_type = \"default\"\n",
    "import sys\n",
    "sys.path.append(\"..\")\n",
    "from segment_anything import sam_model_registry, SamPredictor\n",
    "\n",
    "sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)\n",
    "sam.to(device=device)\n",
    "\n",
    "predictor = SamPredictor(sam)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 第一个无标注数据集的路径\n",
    "# no_label_imgs_path = '/data/yangshuo/DPLABV3P/Ali_building_2class/test_a'\n",
    "# images = glob.glob(no_label_imgs_path + '/*.jpg')  # 读取所有的jpg\n",
    "\n",
    "# 第二个无标注数据集的路径\n",
    "no_label_imgs_path = '/data/yangshuo/DPLABV3P/dfc/img'\n",
    "images = glob.glob(no_label_imgs_path + '/*.tif')  # 读取所有的jpg\n",
    "\n",
    "output_path = '/data/yangshuo/DPLABV3P/sam'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "dbscan = DBSCAN(eps=11, min_samples=300) # 根据实际情况调整参数 eps 和 min_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3578/3578 [42:23<00:00,  1.41it/s]  \n"
     ]
    }
   ],
   "source": [
    "for idx, name in enumerate(tqdm(images)):\n",
    "    \n",
    "    file_name = name.split('/')[-1][:-3] + 'png'\n",
    "    if os.path.exists(osp.join(output_path , file_name )) : continue\n",
    "    \n",
    "    is_exist = True \n",
    "    image = cv2.imread(name)\n",
    "    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    predictor.set_image(image)\n",
    "    \n",
    "    \n",
    "#######################################################################################################\n",
    "    # 原始图片可视化\n",
    "#     plt.figure(figsize=(10,10))\n",
    "#     plt.imshow(image)\n",
    "#     plt.axis('on')\n",
    "#     plt.show()\n",
    "#######################################################################################################\n",
    "    # 获得初始推理\n",
    "    result = inference_segmentor(model, name)  \n",
    "\n",
    "#######################################################################################################\n",
    "    #初始推理可视化\n",
    "#     show_result_pyplot(model, image, result, [[0,0,0],[255,255,255]])\n",
    "    \n",
    "#######################################################################################################\n",
    "    img = np.array(result[0]) \n",
    "    prospect_idx = np.where(img == 1)   # 拿到前景坐标\n",
    "    background_idx = np.where(img == 0 )   # 拿到背景坐标\n",
    "    prospect_coords = np.column_stack((prospect_idx[1], prospect_idx[0]))  # 获取前景的坐标\n",
    "    if len(prospect_coords) == 0 : \n",
    "        is_exist = False  # 如果基础分割器认为当前图片没有建筑\n",
    "        \n",
    "    \n",
    "    background_coords = np.column_stack((background_idx[1], background_idx[0]))  # 获取背景的坐标\n",
    "    \n",
    "    masks = None\n",
    "    if len(background_coords) == 0 : \n",
    "        masks = np.ones([512,512])\n",
    "#     print(background_coords)\n",
    "#     break\n",
    "########################################################################################################\n",
    "    # 下面是prompt策略\n",
    "    \n",
    "    # 对背景随机挑选三十个像素坐标，我们默认所有图片的背景像素个数都在30个以上\n",
    "    indices = 0 \n",
    "    if len(background_coords) != 0 : \n",
    "        indices = np.random.choice(np.arange(len(background_coords)), size=30, replace=True)  \n",
    "        background_coords = background_coords[indices]\n",
    "    \n",
    "    if is_exist:\n",
    "        # 为了增强sam的分割效果，我们对前景的所有坐标随机挑选30个坐标，我们默认如果当前图片有建筑的情况下，像素坐标有30个以上\n",
    "        random_prospect_coords = prospect_coords[np.random.choice(np.arange(len(prospect_coords)), size=30, replace=True)]\n",
    "\n",
    "\n",
    "        # DBSCAN做聚类\n",
    "        dbscan.fit(prospect_coords)  # 获取实例坐标代表\n",
    "        # 获取聚类后的中心点坐标\n",
    "        unique_labels = set(dbscan.labels_)\n",
    "        centers = []\n",
    "        for label in unique_labels:\n",
    "            if label == -1: # 跳过噪声点\n",
    "                continue\n",
    "            class_member_mask = (dbscan.labels_ == label)\n",
    "            center = np.average(prospect_coords[class_member_mask], axis=0)\n",
    "            centers.append(center.astype(int))\n",
    "            # 返回每个聚类中心的坐标\n",
    "            representatives = np.array(centers)\n",
    "            \n",
    "    if representatives.shape == (0,) : # 选不出代表\n",
    "        is_exist = False\n",
    "    ###################################################################################################\n",
    "    # SAM分割阶段\n",
    "    \n",
    "    if is_exist :\n",
    "        # prompt点进行拼接\n",
    "        input_point = np.concatenate((random_prospect_coords , representatives , background_coords) , axis = 0 )\n",
    "\n",
    "        # label拼接\n",
    "        prospect_label = np.ones(len(random_prospect_coords) + len(representatives))  # 前景\n",
    "        background_label = np.zeros(len(background_coords))\n",
    "        input_label = np.append(prospect_label , background_label)\n",
    "\n",
    "\n",
    "        masks, scores, logits = predictor.predict(\n",
    "            point_coords=input_point,\n",
    "            point_labels=input_label,\n",
    "            multimask_output=False\n",
    "        )\n",
    "        mask_input = logits[np.argmax(scores), :, :]  # Choose the model's best mask\n",
    "        masks, _, _ = predictor.predict(\n",
    "            point_coords=input_point,\n",
    "            point_labels=input_label,\n",
    "            mask_input=mask_input[None, :, :],\n",
    "            multimask_output=False\n",
    "        )\n",
    "    else :  # 不存在建筑\n",
    "        # \n",
    "        masks = np.zeros([512,512])\n",
    "        \n",
    "    # 设置布尔数组和保存的文件名\n",
    "    \n",
    "\n",
    "    # 将布尔数组转化为8位深度的单通道像素矩阵，并赋予真值255，假值0\n",
    "    bool_array_8bit = masks.astype(np.uint8) * 255\n",
    "    if bool_array_8bit.shape[0] == 1 : \n",
    "        bool_array_8bit = bool_array_8bit.squeeze(0)\n",
    "    # 将像素矩阵保存为png图片\n",
    "    cv2.imwrite(osp.join(output_path , file_name ), bool_array_8bit)  # 直接是可视化的结果\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
